# -*- coding: utf-8 -*-
# @Time    : 2022/12/10 18:42
# @Author  : Z.J. Zhang
# @Email   : zijingzhang@mail.ustc.edu.cn
# @File    : test.py
# @Software: PyCharm
# 用于求理想磁场下的积分场
import sys
import matplotlib
matplotlib.use('tkagg')
import matplotlib.pyplot as plt
import numpy
from shapely.geometry import LineString

sys.path.append(r'F:\changeworld\caes')
import pygpt.parse_gdf as pg

BIGNUM = 1e10

# gdf = pg.gdftomemory(r"F:\GPTfiles\LGB\integration\LGB2-res\experi_ideaLGB.gdf")[:-1]
gdf = pg.gdftomemory(r"F:\GPTfiles\LGB\integration\experi_ideaLGB.gdf")[:-1]

Bs = list(gdf[0]['d']['Bs'])
Bs.append(0)
zstarts = list(gdf[0]['d']['zstarts'])
zstarts.append(gdf[0]['p']['total_len'])

zstarts.append(BIGNUM)
step = 0
pts = []
pts_piecewise = [list() for i in range(len(Bs))]

for i in range(len(gdf)):
    z = gdf[i]['d']['z']
    x = gdf[i]['d']['x']
    if z > zstarts[step + 1]:
        step += 1
        if step > len(pts_piecewise) - 1:
            break
    pts_piecewise[step].append([x, z])
    pts.append([x, z])

traj = LineString(pts)
plt.plot(*numpy.array(traj.coords).T[::-1])

bounds = [LineString([[-1., zstart], [1., zstart]]) for zstart in zstarts][1:-1]
bound_pts = []
for bound in bounds:
    bound_pts.append(traj.intersection(bound))

for i in range(len(pts_piecewise) - 1):
    bound_pt = list(bound_pts[i].coords)[0]
    pts_piecewise[i].append(bound_pt)
    pts_piecewise[i + 1].insert(0, bound_pt)
lines = [LineString(pts_) for pts_ in pts_piecewise]
step_lens = [line.length for line in lines]
intBs = numpy.array(Bs) * numpy.array(step_lens)
print("积分场大小：%.15f T m\n各台阶积分场：%s T m\nLGB内轨迹全长：%.15f m\n各台阶轨迹长度：%s m" % (
    intBs.sum(), intBs[:-1], sum(step_lens[:-1]), step_lens[:-1]))
base_color = int('1f77b4', 16)
for i in range(len(zstarts) - 2):
    plt.axvspan(zstarts[i], zstarts[i + 1], alpha=0.3, facecolor=[(.0 + 7 * i * 2.1) % 1.0,
                                                                  (.0 + 2 * i * 1.3) % 1.0,
                                                                  (.0 + 3 * i * 1.1) % 1.0])

# plt.axis('equal')
plt.show()
